【DLR】SageMaker Neoでコンパイルしたモデルを動かすときに使えるTips（起動高速化・処理高速化）

カフェチームの山本です。

作成した学習モデルをSageMaker Neoを利用してコンパイルし、Jetson Xavier NXなどのシングルボードで動かす、といった場面があると思います。今回は、コンパイルしたモデルを動かす際に利用するDLR（Deep Learning Runtime）について、使える小ネタを紹介したいと思います。簡単なものですが、起動高速化・処理高速化など、利用するとかなり便利なのでおすすめです。

DLR（Deep Learning Runtime）

シングルボードなどで学習モデルを動かしたいとき、モデルそのままでは大きく動かない（もしくは処理が遅い）という場合があります。そうしたモデルはSageMaker Neoでコンパイルすることで、指定したデバイス用に最適化することができます。このモデルは、DLR（Deep Learning Runtime）というライブラリで動くようになっています。

今回は、このDLRについて、紹介されることが少ないものの、結構便利なTipsを紹介します。

SageMaker Neoでモデルをコンパイル（最適化）して動かす方法自体については、詳しくは同じカフェチームのSINの執筆した記事で詳しく解説されているため、そちらをご覧ください。

[Amazon SageMaker] イメージ分類のモデルをNeoで最適化して、Jetson Nano+OpenCV+Webカメラで使用してみました | Developers.IO

DLRで使える小ネタ

今回紹介するのは主に、以下のページの内容です

Additional Options for TensorRT Optimized Models - dlr documentation

モデルを動かすスクリプトとしては、以下のようなものを用意したとします。ファイル名はexample.pyとします。（このスクリプトである必要はなく、DLRでモデルを利用している一例です）

import dlr
import numpy as np

# Load model.
# /path/to/model is a directory containing the compiled model artifacts (.so, .params, .json)
model = dlr.DLRModel('/path/to/model', 'gpu', 0)

# Prepare some input data.
x = np.random.rand(1, 3, 224, 224)

# Run inference.
y = model.run(x)

通常、スクリプトを起動するコマンドは以下のようだと思います。

python3 example.py

この起動方法を、フラグ（Additional Flags）を追加してから起動することによって、動作を変えられます。フラグは複数同時に使用できます。

ビルドのキャッシュ（起動高速化）

DLRの仕様として、通常の起動方法だと、毎回モデルを読み込む度にビルドし直すようになっています。そのため、毎回起動するたびに待つ必要がありました。（開発中に何度も試行錯誤しようとすると、開発スピードが遅くなり、やる気も低下するという問題がありました。）

起動方法を以下のようにすることで、ビルドした内容をキャッシュし、保存することができます。これによって、1回目の起動時にはビルドが行われるため時間がかかりますが、2回目以降はキャッシュした内容がロードされるだけであるため、時間が大幅に短縮されます。

TVM_TENSORRT_CACHE_DIR=. python3 example.py

指定したフォルダに「9326145774222630086.meta」「9326145774222630086.plan」のようなファイルが出力されます。

具体的には、Jetson Xavier NXにおいて自分が利用してる骨格検出のモデル（SageMaker Neoでコンパイルしたもの）の場合、通常のビルドで2分ほどかかっていた起動が、キャッシュのロードのみでは6秒ほどに短縮されました。かなり開発が効率化されました。おすすめです。

モデルのビット数の変換（処理速度高速化）

モデルを浮動小数点16ビットに変換して実行します。上のページに「これによりパフォーマンスが大幅に向上しますが、モデルの精度が若干低下する可能性があります」と書かれているように、ビット数を落とすことで精度を少し下げながらも、速度を大幅に上げることができます。

一般的に、16bitにしてもそこまで精度が落ちず、むしろ高速化の度合いが大きいことから、特にエッジデバイスで良く利用される方法です。

コマンドとしては、以下のようです。

TVM_TENSORRT_USE_FP16=1 python3 example.py

具体的には、先程の自分が利用しているモデルの場合ですが、処理時間が140[ms]程度から60[ms]程度になり、2倍以上の高速化ができました（処理結果の精度の方は確認できていません）。

SageMaker Neoなどでモデルを再度作成する必要がなく、コマンド1つで処理を高速化できるためかなり便利です。1度は16bitで試しに実行し精度を確認してみて、問題なさそうなら採用する、という流れが良いと思います。

ワークスペースサイズの変更（スループット増加）

モデルを動かす際に利用するGPUのサイズを変更できます。層の深いモデルや、入力サイズが大きいモデルを利用する場合や、データをまとめて入力したいときに利用できるようです。

上のリンクのページに「Jetson NanoおよびJetson TX1ターゲットの場合は最大ワークスペースサイズを256メガバイトに、他のすべてのNVIDIA GPUターゲットの場合は1ギガバイトに自動的に設定します」と書かれているように、SageMaker Neoでコンパイルしたモデルは、動くボードによって自動でこの値を、もともと設定しているようです。この値を上書きして使用します。

コマンドとしては、以下のようです（2GBの場合）。

TVM_TENSORRT_MAX_WORKSPACE_SIZE=2147483647 python3 run.py

自分は今の所使う機会がなく、具体的な効果がわかっていません。

まとめ

今回は、DLRを利用する際に、コマンド1つで起動高速化・処理速度高速化・スループット増加できる方法を紹介しました。

3つのうち最初の2つは、実際に使ってみたところ大きな効果があったのでオススメです。（最後の1つはまだ効果がわかっていません。）

参考にさせていただいたページ

Additional Options for TensorRT Optimized Models - dlr documentation

neo-ai/neo-ai-dlr

[Amazon SageMaker] イメージ分類のモデルをNeoで最適化して、Jetson Nano+OpenCV+Webカメラで使用してみました | Developers.IO